package com.kalessil.phpStorm.phpInspectionsEA.inspectors.security;

import com.intellij.codeInspection.LocalInspectionTool;
import com.intellij.codeInspection.ProblemHighlightType;
import com.intellij.codeInspection.ProblemsHolder;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiElementVisitor;
import com.jetbrains.php.lang.psi.elements.ArrayAccessExpression;
import com.jetbrains.php.lang.psi.elements.ArrayCreationExpression;
import com.jetbrains.php.lang.psi.elements.FunctionReference;
import com.jetbrains.php.lang.psi.elements.Variable;
import com.kalessil.phpStorm.phpInspectionsEA.openApi.BasePhpElementVisitor;
import com.kalessil.phpStorm.phpInspectionsEA.openApi.PhpLanguageLevel;
import com.kalessil.phpStorm.phpInspectionsEA.utils.MessagesPresentationUtil;
import com.kalessil.phpStorm.phpInspectionsEA.utils.OpenapiTypesUtil;
import com.kalessil.phpStorm.phpInspectionsEA.utils.PhpLanguageUtil;
import com.kalessil.phpStorm.phpInspectionsEA.utils.PossibleValuesDiscoveryUtil;
import org.jetbrains.annotations.NotNull;

import java.util.*;

/*
 * This file is part of the Php Inspections (EA Extended) package.
 *
 * (c) Vladimir Reznichenko <kalessil@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

public class UnserializeExploitsInspector extends LocalInspectionTool {
    private static final String messageUseSecondArgument = "Please specify classes allowed for unserialization in 2nd argument.";
    private static final String messagePattern           = "Perhaps it's possible to exploit the unserialize via: %e%.";

    private final static Set<String> untrustedVars      = new HashSet<>();
    private final static Set<String> untrustedFunctions = new HashSet<>();
    static {
        untrustedVars.add("_GET");
        untrustedVars.add("_POST");
        untrustedVars.add("_REQUEST");
        untrustedVars.add("_FILES");
        untrustedVars.add("_COOKIE");

        untrustedFunctions.add("file_get_contents");
        untrustedFunctions.add("base64_decode");
        untrustedFunctions.add("urldecode");
    }

    @NotNull
    @Override
    public String getShortName() {
        return "UnserializeExploitsInspection";
    }

    @NotNull
    @Override
    public String getDisplayName() {
        return "Exploiting unserialize (PHP Object Injection Vulnerability)";
    }

    @Override
    @NotNull
    public PsiElementVisitor buildVisitor(@NotNull final ProblemsHolder holder, boolean isOnTheFly) {
        return new BasePhpElementVisitor() {
            @Override
            public void visitPhpFunctionCall(@NotNull FunctionReference reference) {
                final String functionName = reference.getName();
                if (functionName != null && functionName.equals("unserialize")) {
                    final boolean supportsOptions = PhpLanguageLevel.get(holder.getProject()).atLeast(PhpLanguageLevel.PHP700);
                    final PsiElement[] arguments  = reference.getParameters();
                    if (arguments.length == 1 && !this.isTestContext(reference)) {
                        /* pattern: use 2nd argument since PHP7 */
                        if (supportsOptions) {
                            holder.registerProblem(
                                    reference,
                                    MessagesPresentationUtil.prefixWithEa(messageUseSecondArgument)
                            );
                        }
                        /* pattern: exploitable calls */
                        this.inspectExploits(holder, arguments[0]);
                    } else if (arguments.length == 2 && !this.isTestContext(reference)) {
                        if (arguments[1] instanceof ArrayCreationExpression) {
                            final boolean hasClassesListed = arguments[1].getChildren().length > 0;
                            if (!hasClassesListed) {
                                holder.registerProblem(
                                        reference,
                                        MessagesPresentationUtil.prefixWithEa(messageUseSecondArgument)
                                );
                            }
                        } else if (PhpLanguageUtil.isTrue(arguments[1])) {
                            holder.registerProblem(
                                    reference,
                                    MessagesPresentationUtil.prefixWithEa(messageUseSecondArgument)
                            );
                        }
                    }
                }
            }

            private void inspectExploits(@NotNull ProblemsHolder holder, @NotNull PsiElement argument) {
                final Set<PsiElement> values = PossibleValuesDiscoveryUtil.discover(argument);
                if (! values.isEmpty()) {
                    final List<String> reporting = new ArrayList<>();
                    for (PsiElement value : values) {
                        if (OpenapiTypesUtil.isFunctionReference(value)) {
                            final FunctionReference call = (FunctionReference) value;
                            final String functionName    = call.getName();
                            if (functionName != null && untrustedFunctions.contains(functionName)) {
                                reporting.add(functionName + "(...)");
                            }
                            continue;
                        }

                        /* extract array access variable */
                        if (value instanceof ArrayAccessExpression) {
                            PsiElement container = value;
                            while (container instanceof ArrayAccessExpression) {
                                container = ((ArrayAccessExpression) container).getValue();
                            }
                            if (container instanceof Variable) {
                                value = container;
                            }
                        }

                        if (value instanceof Variable && untrustedVars.contains(((Variable) value).getName())) {
                            reporting.add(value.getText());
                            // continue
                        }

                        /* other expressions are not supported currently */
                    }

                    /* got something for reporting */
                    if (! reporting.isEmpty()) {
                        /* sort reporting list to produce testable results */
                        Collections.sort(reporting);
                        holder.registerProblem(
                                argument,
                                MessagesPresentationUtil.prefixWithEa(messagePattern.replace("%e%", String.join(", ", reporting))),
                                ProblemHighlightType.GENERIC_ERROR
                        );
                        reporting.clear();
                    }
                }
                values.clear();
            }
        };
    }
}
